import numpy as np
from osgeo import gdal, osr
from scipy.spatial.distance import cdist
from scipy.ndimage import binary_erosion, generate_binary_structure
import matplotlib.pyplot as plt
from tqdm import tqdm  # 导入tqdm库
import time  # 导入time模块
import os  # 导入os模块


#  ===============================================================
#  计算景观结构指数 或 景观稳定度指数
#  ===============================================================
class LandscapeMetrics:
    def __init__(self, image_path, window_scale=20, weights=[1 / 3, 1 / 3, 1 / 3], use_padding=True, output_dir=None):
        """
        初始化LandscapeMetrics类, 用于计算景观稳定度指数或称为景观结构指数。

        参数:
        image_path (str): 输入GeoTIFF图像的路径。
        window_scale (int): 窗口大小为图像大小的1/window_scale。
        weights: 破碎度指数的权重、分离度指数的权重、分维数的权重，默认为均为1/3。
        use_padding (bool): 是否启用边缘填充，默认为True。
        output_dir (str): 图形输出目录，默认为None（当前目录）。
        """
        self.image_path = image_path
        self.window_scale = window_scale
        self.a = weights[0]
        self.b = weights[1]
        self.c = weights[2]
        self.use_padding = use_padding
        self.output_dir = output_dir if output_dir else os.getcwd()  # 设置输出目录
        self.data = None
        self.resolution = None
        self.geotransform = None
        self.projection = None
        self.window_size = None
        self.padded_data = None
        self.nodata = None  # 存储输入图像的nodata值
        self.load_data()

    def load_data(self):
        """使用GDAL加载GeoTIFF数据并预处理。"""
        dataset = gdal.Open(self.image_path)
        if dataset is None:
            raise ValueError("无法打开文件")

        # 读取第一个波段的数据，并转换为浮点数
        band = dataset.GetRasterBand(1)
        self.data = band.ReadAsArray().astype(float)

        # 获取空间分辨率、地理变换和投影信息
        self.geotransform = dataset.GetGeoTransform()
        self.resolution = self.geotransform[1]  # 假设为正方形像元
        self.projection = dataset.GetProjection()

        # 获取NoData值
        self.nodata = band.GetNoDataValue()
        if self.nodata is not None:
            self.data[self.data == self.nodata] = np.nan

        # 计算窗口大小
        rows, cols = self.data.shape
        self.window_size = (rows // self.window_scale, cols // self.window_scale)

        # 对输入图像进行边缘填充（如果启用）
        if self.use_padding:
            pad_rows = self.window_size[0] // 2
            pad_cols = self.window_size[1] // 2
            self.padded_data = np.pad(self.data, ((pad_rows, pad_rows), (pad_cols, pad_cols)), mode='constant',
                                      constant_values=np.nan)
        else:
            self.padded_data = self.data  # 如果不启用填充，直接使用原始数据

        dataset = None  # 关闭数据集

    def normalize(self, array):
        """
        将数组归一化到[0, 1]范围，并考虑nodata值的影响。

        参数:
        array (numpy.ndarray): 输入的景观指数数组。

        返回:
        normalized_array (numpy.ndarray): 归一化后的数组。
        """
        # 将nodata值替换为NaN
        if self.nodata is not None:
            array = np.where(array == self.nodata, np.nan, array)

        # 计算最小值和最大值（忽略NaN值）
        array_min = np.nanmin(array)
        array_max = np.nanmax(array)

        # 如果所有值相同，则返回0
        if array_max == array_min:
            return np.zeros_like(array)

        # 归一化到[0, 1]范围
        normalized_array = (array - array_min) / (array_max - array_min)

        # 将NaN值恢复为nodata值
        if self.nodata is not None:
            normalized_array = np.where(np.isnan(normalized_array), self.nodata, normalized_array)

        return normalized_array

    def calculate_fragmentation_index(self, window):
        """计算景观破碎度指数（FI）。"""
        valid_pixels = window[~np.isnan(window)]  # 忽略nodata值
        if len(valid_pixels) == 0:
            return np.nan
        unique_values = np.unique(valid_pixels)
        total_area = len(valid_pixels) * (self.resolution ** 2)
        patch_areas = [np.sum(valid_pixels == value) * (self.resolution ** 2) for value in unique_values]
        min_patch_area = min(patch_areas)
        N_c = total_area / min_patch_area
        FI = (len(unique_values) - 1) / N_c
        return FI

    def calculate_separation_index(self, window):
        """计算景观分离度指数（SI）。"""
        valid_pixels = window[~np.isnan(window)]  # 忽略nodata值
        if len(valid_pixels) == 0:
            return np.nan
        unique_values = np.unique(valid_pixels)
        separation_indices = []
        for value in unique_values:
            patches = np.argwhere(window == value)
            N = len(patches)
            A = np.sum(window == value) * (self.resolution ** 2)
            if N > 1:
                distances = cdist(patches, patches)
                np.fill_diagonal(distances, np.inf)  # 忽略自身
                min_distances = np.min(distances, axis=1)
                SI = 0.5 * np.sqrt(N / A) * np.mean(min_distances) * self.resolution
            else:
                SI = 0  # 如果只有一个斑块，分离度为0
            separation_indices.append(SI)
        return np.mean(separation_indices)

    def calculate_fractal_dimension_index(self, window):
        """计算景观分维度指数（FDI）。"""
        valid_pixels = window[~np.isnan(window)]  # 忽略nodata值
        if len(valid_pixels) == 0:
            return np.nan
        unique_values = np.unique(valid_pixels)
        fractal_dimensions = []
        for value in unique_values:
            binary_patch = (window == value).astype(int)
            perimeter = self.calculate_perimeter(binary_patch)
            area = np.sum(binary_patch) * (self.resolution ** 2)
            if area > 0 and perimeter > 0:
                FDI = 2 * np.log(perimeter) / np.log(area)
            else:
                FDI = 0
            fractal_dimensions.append(FDI)
        return np.mean(fractal_dimensions)

    def calculate_patch_density_index(self, window):
        """计算斑块密度指数（PD）。"""
        valid_pixels = window[~np.isnan(window)]  # 忽略nodata值
        if len(valid_pixels) == 0:
            return np.nan
        unique_values = np.unique(valid_pixels)
        total_area = len(valid_pixels) * (self.resolution ** 2)
        PD = len(unique_values) / total_area
        return PD

    def calculate_perimeter(self, binary_patch):
        """计算斑块的周长。"""
        struct = generate_binary_structure(2, 2)
        eroded = binary_erosion(binary_patch, structure=struct)
        edge = binary_patch ^ eroded
        return np.sum(edge) * self.resolution

    def calculate_landscape_stability(self, fi, si, fdi):
        """
        计算景观稳定性指数 (LS)。

        参数:
        fi (float): 归一化的破碎度指数。
        si (float): 归一化的分离度指数。
        fdi (float): 归一化的分维度指数。

        返回:
        LS (float): 景观稳定性指数。
        """
        LS = 1 - (self.a * fi + self.b * si + self.c * fdi)
        return LS

    def calculate_patch_density_index(self, window):
        """计算斑块密度指数（PD）。"""
        unique_values = np.unique(window[~np.isnan(window)])
        if len(unique_values) == 0:
            return np.nan
        total_area = window.size * (self.resolution ** 2)
        PD = len(unique_values) / total_area
        return PD

    def calculate_window_metrics(self):
        """使用移动窗口法计算景观指数，并保存为GeoTIFF文件。"""
        rows, cols = self.data.shape
        window_rows, window_cols = self.window_size
        pad_rows = window_rows // 2 if self.use_padding else 0
        pad_cols = window_cols // 2 if self.use_padding else 0

        # 初始化输出数组
        fi_array = np.zeros((rows, cols), dtype=np.float32)
        si_array = np.zeros_like(fi_array)
        fdi_array = np.zeros_like(fi_array)
        pd_array = np.zeros_like(fi_array)  # 计算斑块密度指数PD
        # ls_array = np.zeros_like(fi_array)

        # 计时开始
        start_time = time.time()

        # 滑动窗口计算景观指数（使用tqdm显示进度条）
        for i in tqdm(range(pad_rows, rows + pad_rows), desc="Processing rows"):
            for j in range(pad_cols, cols + pad_cols):
                window = self.padded_data[i - pad_rows:i + pad_rows + 1, j - pad_cols:j + pad_cols + 1]
                fi = self.calculate_fragmentation_index(window)
                si = self.calculate_separation_index(window)
                fdi = self.calculate_fractal_dimension_index(window)
                pd = self.calculate_patch_density_index(window)  # 计算斑块密度指数PD
                fi_array[i - pad_rows, j - pad_cols] = fi
                si_array[i - pad_rows, j - pad_cols] = si
                fdi_array[i - pad_rows, j - pad_cols] = fdi
                pd_array[i - pad_rows, j - pad_cols] = pd  # 计算斑块密度指数PD

        # 计时结束
        end_time = time.time()
        print(f"计算完成，耗时: {end_time - start_time:.2f} 秒")

        # 归一化景观指数
        fi_normalized = self.normalize(fi_array)
        si_normalized = self.normalize(si_array)
        fdi_normalized = self.normalize(fdi_array)
        pd_normalized = self.normalize(pd_array)  # 计算斑块密度指数PD

        # 计算景观稳定性指数 (LS)
        ls_array = self.calculate_landscape_stability(fi_normalized, si_normalized, fdi_normalized)
        ls_normalized = self.normalize(ls_array)

        # 保存结果为GeoTIFF文件
        self.save_as_geotiff(fi_normalized, "fragmentation_index_normalized.tif")
        self.save_as_geotiff(si_normalized, "separation_index_normalized.tif")
        self.save_as_geotiff(fdi_normalized, "fractal_dimension_index_normalized.tif")
        self.save_as_geotiff(pd_normalized, "patch_density_index_normalized.tif")  # 计算斑块密度指数PD
        self.save_as_geotiff(ls_array, "landscape_stability_index.tif")
        self.save_as_geotiff(ls_normalized, "landscape_stability_index_normalized.tif")

        # 可视化结果
        self.visualize_results(fi_normalized, "Fragmentation Index (Normalized)", "fragmentation_index_normalized.png")
        self.visualize_results(si_normalized, "Separation Index (Normalized)", "separation_index_normalized.png")
        self.visualize_results(fdi_normalized, "Fractal Dimension Index (Normalized)",
                               "fractal_dimension_index_normalized.png")
        self.visualize_results(pd_normalized, "Patch Density Index (Normalized)",
                               "patch_density_index_normalized.png")  # 计算斑块密度指数PD
        self.visualize_results(ls_array, "Landscape Stability Index",
                               "landscape_stability_index.png")
        self.visualize_results(ls_normalized, "Landscape Stability Index (Normalized)",
                               "landscape_stability_index_normalized.png")

    def save_as_geotiff(self, array, output_filename):
        """将数组保存为GeoTIFF文件，并应用输入图像的nodata范围。"""
        rows, cols = array.shape
        driver = gdal.GetDriverByName("GTiff")
        output_path = os.path.join(self.output_dir, output_filename)  # 拼接输出路径
        out_dataset = driver.Create(output_path, cols, rows, 1, gdal.GDT_Float32)
        out_dataset.SetGeoTransform(self.geotransform)
        out_dataset.SetProjection(self.projection)
        out_band = out_dataset.GetRasterBand(1)

        # 将输入图像的nodata范围应用到输出图像
        if self.nodata is not None:
            array[np.isnan(array)] = self.nodata  # 将NaN替换为nodata值
            out_band.SetNoDataValue(self.nodata)

        out_band.WriteArray(array)
        out_band.FlushCache()
        out_dataset = None

    def visualize_results(self, array, title, output_filename):
        """可视化景观指数结果并保存为PNG文件。"""
        output_path = os.path.join(self.output_dir, output_filename)  # 拼接输出路径
        plt.figure(figsize=(10, 10))
        plt.imshow(array, cmap='viridis', vmin=0, vmax=1)
        plt.colorbar(label=title)
        plt.title(title)
        plt.axis('off')
        plt.savefig(output_path, bbox_inches='tight', dpi=300)
        plt.close()


# 示例使用
if __name__ == "__main__":
    # =============================1.计算景观结构指数===============================
    image_path = "F:/test/lu2022_1k.tif"  # 输入GeoTIFF图像路径
    output_path = "F:/test/LS"  # 输出相关文件路径
    weights = [0.3, 0.4, 0.3]  # 各个景观格局指标的权重，其权重采用层次分析法计算，可自行赋值，默认定义为0.3、0.4、0.3。

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # 创建LandscapeMetrics对象
    landscape = LandscapeMetrics(
        image_path,
        window_scale=20,
        weights=weights,
        use_padding=True,  # 启用边缘填充
        output_dir=output_path  # 设置图形输出目录
    )

    # 计算窗口景观指数并保存为GeoTIFF文件
    landscape.calculate_window_metrics()
